# -*- coding: utf-8 -*-
"""
-------------------------------------------------
   File Name：     num_detector
   Description :
   Author :       'li'
   date：          2021/7/4
-------------------------------------------------
   Change Activity:
                   2021/7/4:
-------------------------------------------------
"""
from ml.cv.algorithms.bbox_generate.corner_to_bbox.corner_to_bbox_generator import CornerToBboxGenerator
from ml.cv.algorithms.bbox_generate.corner_to_bbox.format_gap_mask import GapMaskFormat
from ml.modules.container_num_recognition.single_container_recognition.module.check_result_template import clip_text_area


class NumDetector:
    def __init__(self, img=None, torch_detector=None):
        self.np_img = img
        self.torch_detector = torch_detector

    def get_bboxes_areas(self):
        textarea_mask, gap_mask = self.torch_detector.get_text_gap_mask(self.np_img)
        gap_mask = GapMaskFormat(gap_mask).format_gap_mask()
        if gap_mask is not None:
            textarea_mask = textarea_mask - gap_mask
        bboxes = CornerToBboxGenerator(textarea_mask).gen_bboxes()
        results = []
        for bbox in bboxes:
            area, bbox = clip_text_area(bbox, self.np_img)
            if area is None:
                continue
            results.append([area, bbox])
        return results
